import os, shutil
import os.path as osp
import json
import numpy as np
from tqdm import tqdm
from tool import filesystem, via_tool # export PYTHONPATH=$PYTHONPATH:`pwd`

def txt_to_via(image_dir, 
                save_dir, 
                label_dir,
                sub_dir="ua_detrac", 
                label_gt="train_gt.txt",
                label_ign="train_ign.txt",
                via_name="via_region_data.json"):
    """
    ua detrac tool 
    将此数据集转换成via格式
    """
    data_dict = dict()
    ignore_dict = dict()
    with open(osp.join(label_dir, label_gt), "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            data_dict[info[0]] = info[1:]
    with open(osp.join(label_dir, label_ign), "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            ignore_dict[info[0]] = info[1:]

    for sub_name in tqdm(os.listdir(image_dir)):
        cur_dir = osp.join(image_dir, sub_name)
        cur_save_dir = osp.join(save_dir, sub_dir, sub_name)
        os.makedirs(cur_save_dir, exist_ok=True)

        total_json = dict()
        for file_name in os.listdir(cur_dir):
            boxs = data_dict.get(osp.join(sub_name, file_name), [])
            if boxs:
                boxs = np.array(boxs, dtype=np.float32).reshape(-1, 4)

            file_path = osp.join(cur_dir, file_name)
            file_size = osp.getsize(file_path)
            one_json = dict()
            one_json["filename"] = file_name
            one_json["size"] = file_size
            one_json["file_attributes"] = {}
            
            regions = []
            for box in boxs:
                region = dict()
                shape_attributes = {
                    "name": "rect",
                    "x": int(box[0]),
                    "y": int(box[1]),
                    "width": int(box[2] - box[0]),
                    "height": int(box[3] - box[1])
                }
                region_attributes = { "label": "car" }
                region["shape_attributes"] = shape_attributes
                region["region_attributes"] = region_attributes
                regions.append(region)
            one_json["regions"] = regions
            total_json[file_name + str(file_size)] = one_json
            shutil.copy(file_path, osp.join(cur_save_dir, file_name))

        with open(cur_save_dir + os.sep + via_name, "w") as wf:
            wf.write(json.dumps(total_json))



if __name__ == "__main__":
    # image_dir = "/home/xc/work/data/car/UA-DETRAC/Insight-MVT_Annotation_Train"
    # save_dir = "/home/xc/work/data/car/train"
    # label_dir = "/home/xc/work/data/car/UA-DETRAC"
    # txt_to_via(image_dir, save_dir, label_dir)

